package common

import (
	"context"
	"sync"
	"time"
)

// AsyncTaskManager 异步任务管理器
type AsyncTaskManager struct {
	workers    chan struct{}
	ctx        context.Context
	cancel     context.CancelFunc
	wg         sync.WaitGroup
	maxWorkers int
}

// NewAsyncTaskManager 创建异步任务管理器
func NewAsyncTaskManager(maxWorkers int) *AsyncTaskManager {
	ctx, cancel := context.WithCancel(context.Background())
	return &AsyncTaskManager{
		workers:    make(chan struct{}, maxWorkers),
		ctx:        ctx,
		cancel:     cancel,
		maxWorkers: maxWorkers,
	}
}

// Execute 执行异步任务
func (tm *AsyncTaskManager) Execute(task func()) error {
	select {
	case tm.workers <- struct{}{}: // 获取工作槽
		tm.wg.Add(1)
		go func() {
			defer func() {
				<-tm.workers // 释放工作槽
				tm.wg.Done()
			}()
			task()
		}()
		return nil
	case <-tm.ctx.Done():
		return tm.ctx.Err()
	default:
		// 如果没有可用的工作槽，直接同步执行
		task()
		return nil
	}
}

// Shutdown 关闭任务管理器
func (tm *AsyncTaskManager) Shutdown(timeout time.Duration) {
	tm.cancel()
	done := make(chan struct{})
	go func() {
		tm.wg.Wait()
		close(done)
	}()

	select {
	case <-done:
		return
	case <-time.After(timeout):
		return
	}
}
